import math
import os

import tables

import charmicat
import histogram
from fitter import GaussPlusPolynomialFit

import cmix.cutdef
import cmix.histdef
import cmix.parameters

def fom_count(cutset, rs_files, ws_files, denom_cut='m_kpi > -1e9'):
    if type(cutset) == type("hello!"):
        if cutset in cmix.cutdef.cutsets:
            cutset = [cmix.cutdef.cuts[n] for n in cmix.cutdef.cutsets[cutset]]
        else:
            cutset = [charmicat.Cut(name=cutset,
                                    desc=cutset,
                                    pos_str=cutset)]

    elif type(cutset) == charmicat.event_filter.Cut:
        cutset = [cutset]
    else:
        msg = "Don't know how to handle cutset of type `{0}'"
        print msg.format(type(cutset))
        return

    cut = charmicat.cut_intersection(cutset)

    if type(rs_files) == type("hello!"):
        rs_files = [rs_files]

    if type(ws_files) == type("hello!"):
        ws_files = [ws_files]

    rsdenom = 0
    rsnum   = 0
    for rsf in rs_files:
        with tables.openFile(rsf) as fp:
            rsdenom += len(fp.root.nt.getWhereList(denom_cut,
                           condvars=cmix.parameters.paramdef))
            rsnum += len(fp.root.nt.getWhereList(cut.pos_str,
                                                 condvars=cmix.parameters.paramdef))
    
    rseff = float(rsnum) / float(rsdenom)
        
    wsdenom = 0
    wsnum   = 0
    for wsf in ws_files:
        with tables.openFile(wsf) as fp:
            wsdenom += len(fp.root.nt.getWhereList(denom_cut,
                                                   condvars=cmix.parameters.paramdef))
            wsnum += len(fp.root.nt.getWhereList(cut.pos_str,
                                                 condvars=cmix.parameters.paramdef))
                    
    wseff = float(wsnum) / float(wsdenom)

    return wseff / math.sqrt((wseff + 300*rseff))
    


def fom_fit(cutset, rs_files, ws_files, hist, dbg={}, denom_cut='m_kpi > -1e9', 
            polyorder=4, plot=None):

    if type(cutset) == type("hello!"):
        if cutset in cmix.cutdef.cutsets:
            cutset = [cmix.cutdef.cuts[n] for n in cmix.cutdef.cutsets[cutset]]
        else:
            cutset = [charmicat.Cut(name=cutset,
                                    desc=cutset,
                                    pos_str=cutset)]
        
    elif type(cutset) == charmicat.Cut:
        cutset = [cutset]

    else:
        msg = "Don't know how to handle cutset of type `{0}'"
        print msg.format(type(cutset))
        return
    
    if type(hist) == type("hello!"):
        hist = cmix.histdef.hists[hist]
    elif type(hist) in histogram.htypes:
        pass
    else:
        msg = "Don't know how to handle histogram of type `{0}'"
        print msg.format(type(cutset))
        return

    cut = charmicat.cut_intersection(cutset)

    if type(rs_files) == type("hello!"):
        rs_files = [rs_files]

    if type(ws_files) == type("hello!"):
        ws_files = [ws_files]

    rsdenom = 0
    rsnum_hist = histogram.copy(hist)

    for rsf in rs_files:
        with tables.openFile(rsf) as fp:            
            rsdenom += len(fp.root.nt.getWhereList(denom_cut, condvars=cmix.params))
            rsnum_hist.fill(fp.root.nt.readWhere(cut.pos_str, field='m_kpi', 
                                                 condvars=cut.condvars))
            
            
    wsdenom = 0
    wsnum_hist = histogram.copy(hist)
            
    for wsf in ws_files:
        with tables.openFile(wsf) as fp:          
            wsdenom += len(fp.root.nt.getWhereList(denom_cut, condvars=cmix.params))
            wsnum_hist.fill(fp.root.nt.readWhere(cut.pos_str, field='m_kpi', 
                                                 condvars=cut.condvars))


    frs = GaussPlusPolynomialFit(rsnum_hist, polyorder)

    frs.fix_signal()
    frs.m.printMode = 0
    frs.m.fixed['A'] = False
    frs.m.values['mu'] = cmix.parameters.paramdef['D0_MASS']
    frs.m.values['sigma'] = cmix.parameters.paramdef['D0_SIGMA_MC']

    frs.fit(*rsnum_hist.arange)

    if plot is not None:
        bn = os.path.basename(plot).replace('_', ' ')
        frs.plot(name='{0}_rsfit'.format(plot), title='{0} RS'.format(bn))

    fws = GaussPlusPolynomialFit(wsnum_hist, polyorder)

    fws.fix_signal()
    fws.m.printMode = 0
    fws.m.fixed['A'] = False
    fws.m.values['mu'] = cmix.parameters.paramdef['D0_MASS']
    fws.m.values['sigma'] = cmix.parameters.paramdef['D0_SIGMA_MC']

    fws.fit(*wsnum_hist.arange)

    if plot is not None:
        bn = os.path.basename(plot).replace('_', ' ')
        fws.plot(name='{0}_wsfit'.format(plot), title='{0} WS'.format(bn))

    eff_rs = frs.m.values['A'] / rsdenom
    eff_ws = fws.m.values['A'] / wsdenom
    
    dbg['rsdenom'] = rsdenom
    dbg['wsdenom'] = wsdenom
    dbg['rsA'] = frs.m.values['A']
    dbg['wsA'] = fws.m.values['A']

    return eff_ws / math.sqrt(eff_ws + 300*eff_rs)

                                   
